{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f5f09249",
   "metadata": {},
   "source": [
    "# Agents using RAG via LlamaIndex\n",
    "\n",
    "This notebook highlights how to use the LlamaIndex tools in **Docling MCP**.\n",
    "\n",
    "We will use the Llama Stack backend to run a local Responses API system.\n",
    "\n",
    "### Tools:\n",
    "\n",
    "We will use tools from the Docling MCP server that allow executing tasks such as:\n",
    "- converting a PDF file from a local or remote location into a unified document representation [DoclingDocument](https://docling-project.github.io/docling/concepts/docling_document/).\n",
    "- chunk and ingest the document in the LlamaIndex vectordb.\n",
    "- search in the document using rag techniques.\n",
    "\n",
    "### Runtime:\n",
    "\n",
    "- backend: use LLama Stack providing an OpenAI-compatible Responses API\n",
    "- client: use the OpenAI Agents SDK"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6621973b",
   "metadata": {},
   "source": [
    "## Pre-Requisites\n",
    "\n",
    "Before starting this notebook, ensure that you have:\n",
    "1. Followed the instructions in the [Llama Stack README](../llama-stack/README.md) to set up the following resources:\n",
    "  - Inference model with Ollama\n",
    "  - Llama Stack server with the starter template [distribution-starter](https://hub.docker.com/r/llamastack/distribution-starter)\n",
    "\n",
    "2. Started the Docling MCP server with the `conversion` and `llama-index-rag` groups. See the details in the [README](./README.md).\n",
    "\n",
    "You may want to create a virtual environment to run this notebook, for instance, with [uv](https://docs.astral.sh/uv/):\n",
    "\n",
    "```bash\n",
    "uv venv\n",
    "source .venv/bin/activate\n",
    "uv pip install openai-agents rich\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c30ae37e",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e6ec217b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import uuid\n",
    "\n",
    "from rich.console import Console\n",
    "\n",
    "from agents import Agent, ModelSettings, Runner, SQLiteSession, set_trace_processors\n",
    "from agents.mcp import MCPServerStreamableHttp\n",
    "from agents.models.openai_provider import OpenAIProvider\n",
    "from agents.run import RunConfig\n",
    "from agents.tracing.processors import BatchTraceProcessor, ConsoleSpanExporter\n",
    "from openai import AsyncOpenAI\n",
    "\n",
    "console = Console()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e49dbcc2",
   "metadata": {},
   "outputs": [],
   "source": [
    "BASE_URL = \"http://localhost:8321/v1/openai/v1\"\n",
    "API_KEY = \"none\"\n",
    "MODEL_ID = \"llama3.2:3b-instruct-fp16\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "98715ebc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Created session_id=9028ac0e-bc71-47b2-8f97-816ff1194445 for Agent(Doc QA)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #008080; text-decoration-color: #008080\">User&gt; Convert the PDF document on </span><span style=\"color: #008080; text-decoration-color: #008080; text-decoration: underline\">https://arxiv.org/pdf/2206.01062</span><span style=\"color: #008080; text-decoration-color: #008080\"> to DoclingDocument.</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[36mUser> Convert the PDF document on \u001b[0m\u001b[4;36mhttps://arxiv.org/pdf/2206.01062\u001b[0m\u001b[36m to DoclingDocument.\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Exporter] Export trace_id=trace_b77695865cde43b680b34dac44164fb4, name=Agent workflow\n",
      "[Exporter] Export span: {'object': 'trace.span', 'id': 'span_5daf58a2405f4f35b2ddc696', 'trace_id': 'trace_b77695865cde43b680b34dac44164fb4', 'parent_id': None, 'started_at': '2025-08-18T11:00:20.050557+00:00', 'ended_at': '2025-08-18T11:00:20.061691+00:00', 'span_data': {'type': 'mcp_tools', 'server': 'Docling MCP', 'result': ['is_document_in_local_cache', 'convert_document_into_docling_document', 'export_docling_document_to_vector_db', 'search_documents']}, 'error': None}\n",
      "[Exporter] Export span: {'object': 'trace.span', 'id': 'span_af58cd701f914cf3a04fa5a5', 'trace_id': 'trace_b77695865cde43b680b34dac44164fb4', 'parent_id': 'span_2c4f5ae3008c4f608da84b7e', 'started_at': '2025-08-18T11:00:20.062017+00:00', 'ended_at': '2025-08-18T11:00:24.823131+00:00', 'span_data': {'type': 'response', 'response_id': 'resp-e9e3d8f7-57b7-4c2a-9dca-8cb030c7238d'}, 'error': None}\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #008000; text-decoration-color: #008000\">Assistant&gt; The PDF document from the specified URL has been successfully converted into a DoclingDocument. The </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">document key for this conversion is </span><span style=\"color: #008000; text-decoration-color: #008000\">\"868f49ae1f0e66e82238a8aea43fd30b\"</span><span style=\"color: #008000; text-decoration-color: #008000\">.</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[32mAssistant> The PDF document from the specified URL has been successfully converted into a DoclingDocument. The \u001b[0m\n",
       "\u001b[32mdocument key for this conversion is \u001b[0m\u001b[32m\"868f49ae1f0e66e82238a8aea43fd30b\"\u001b[0m\u001b[32m.\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #008080; text-decoration-color: #008080\">User&gt; Export the document to the vectordb.</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[36mUser> Export the document to the vectordb.\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Exporter] Export span: {'object': 'trace.span', 'id': 'span_e6ed81424df14bf9949fb307', 'trace_id': 'trace_b77695865cde43b680b34dac44164fb4', 'parent_id': 'span_2c4f5ae3008c4f608da84b7e', 'started_at': '2025-08-18T11:00:24.823338+00:00', 'ended_at': '2025-08-18T11:00:46.321972+00:00', 'span_data': {'type': 'function', 'name': 'convert_document_into_docling_document', 'input': '{\"source\":\"https://arxiv.org/pdf/2206.01062\"}', 'output': '{\"type\":\"text\",\"text\":\"{\\\\n  \\\\\"success\\\\\": true,\\\\n  \\\\\"document_key\\\\\": \\\\\"868f49ae1f0e66e82238a8aea43fd30b\\\\\"\\\\n}\",\"annotations\":null,\"meta\":null}', 'mcp_data': {'server': 'Docling MCP'}}, 'error': None}\n",
      "[Exporter] Export span: {'object': 'trace.span', 'id': 'span_11ec2106a80249d98f5a8720', 'trace_id': 'trace_b77695865cde43b680b34dac44164fb4', 'parent_id': 'span_2c4f5ae3008c4f608da84b7e', 'started_at': '2025-08-18T11:00:46.322273+00:00', 'ended_at': '2025-08-18T11:00:46.327432+00:00', 'span_data': {'type': 'mcp_tools', 'server': 'Docling MCP', 'result': ['is_document_in_local_cache', 'convert_document_into_docling_document', 'export_docling_document_to_vector_db', 'search_documents']}, 'error': None}\n",
      "[Exporter] Export span: {'object': 'trace.span', 'id': 'span_c853c9016da449b4bf2fb322', 'trace_id': 'trace_b77695865cde43b680b34dac44164fb4', 'parent_id': 'span_2c4f5ae3008c4f608da84b7e', 'started_at': '2025-08-18T11:00:46.328014+00:00', 'ended_at': '2025-08-18T11:00:47.800880+00:00', 'span_data': {'type': 'response', 'response_id': 'resp-a9acdab9-60dc-4bdc-9281-7fb028b194de'}, 'error': None}\n",
      "[Exporter] Export span: {'object': 'trace.span', 'id': 'span_2c4f5ae3008c4f608da84b7e', 'trace_id': 'trace_b77695865cde43b680b34dac44164fb4', 'parent_id': None, 'started_at': '2025-08-18T11:00:20.061769+00:00', 'ended_at': '2025-08-18T11:00:47.802289+00:00', 'span_data': {'type': 'agent', 'name': 'Doc QA', 'handoffs': [], 'tools': ['is_document_in_local_cache', 'convert_document_into_docling_document', 'export_docling_document_to_vector_db', 'search_documents'], 'output_type': 'str'}, 'error': None}\n",
      "[Exporter] Export trace_id=trace_04e821a8ac0744a083a1549fdd6c4465, name=Agent workflow\n",
      "[Exporter] Export span: {'object': 'trace.span', 'id': 'span_fa36e6dbd1e5497999af61ac', 'trace_id': 'trace_04e821a8ac0744a083a1549fdd6c4465', 'parent_id': None, 'started_at': '2025-08-18T11:00:47.807369+00:00', 'ended_at': '2025-08-18T11:00:47.819711+00:00', 'span_data': {'type': 'mcp_tools', 'server': 'Docling MCP', 'result': ['is_document_in_local_cache', 'convert_document_into_docling_document', 'export_docling_document_to_vector_db', 'search_documents']}, 'error': None}\n",
      "[Exporter] Export span: {'object': 'trace.span', 'id': 'span_15f2c14cd5d6410692ca201d', 'trace_id': 'trace_04e821a8ac0744a083a1549fdd6c4465', 'parent_id': 'span_0957dcf269204462a27f7a2d', 'started_at': '2025-08-18T11:00:47.820424+00:00', 'ended_at': '2025-08-18T11:00:49.055879+00:00', 'span_data': {'type': 'response', 'response_id': 'resp-124f2116-5c6f-493f-aad0-6ae6c1c015be'}, 'error': None}\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #008000; text-decoration-color: #008000\">Assistant&gt; The DoclingDocument has been successfully exported to the vector database. The document was initially </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">initialised in the database. If you need more information, feel free to ask!</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[32mAssistant> The DoclingDocument has been successfully exported to the vector database. The document was initially \u001b[0m\n",
       "\u001b[32minitialised in the database. If you need more information, feel free to ask!\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #008080; text-decoration-color: #008080\">User&gt; Search in the document: How many pages were manually annotated in the dataset?</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[36mUser> Search in the document: How many pages were manually annotated in the dataset?\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Exporter] Export span: {'object': 'trace.span', 'id': 'span_590054d85eb8419c92f29eab', 'trace_id': 'trace_04e821a8ac0744a083a1549fdd6c4465', 'parent_id': 'span_0957dcf269204462a27f7a2d', 'started_at': '2025-08-18T11:00:49.056869+00:00', 'ended_at': '2025-08-18T11:00:51.477845+00:00', 'span_data': {'type': 'function', 'name': 'export_docling_document_to_vector_db', 'input': '{\"document_key\":\"868f49ae1f0e66e82238a8aea43fd30b\"}', 'output': '{\"type\":\"text\",\"text\":\"Successful initialisation for document with id 868f49ae1f0e66e82238a8aea43fd30b\",\"annotations\":null,\"meta\":null}', 'mcp_data': {'server': 'Docling MCP'}}, 'error': None}\n",
      "[Exporter] Export span: {'object': 'trace.span', 'id': 'span_c545efc32dd845a38d365fc7', 'trace_id': 'trace_04e821a8ac0744a083a1549fdd6c4465', 'parent_id': 'span_0957dcf269204462a27f7a2d', 'started_at': '2025-08-18T11:00:51.478065+00:00', 'ended_at': '2025-08-18T11:00:51.481436+00:00', 'span_data': {'type': 'mcp_tools', 'server': 'Docling MCP', 'result': ['is_document_in_local_cache', 'convert_document_into_docling_document', 'export_docling_document_to_vector_db', 'search_documents']}, 'error': None}\n",
      "[Exporter] Export span: {'object': 'trace.span', 'id': 'span_d9ddfff85cc947fdb24e0c78', 'trace_id': 'trace_04e821a8ac0744a083a1549fdd6c4465', 'parent_id': 'span_0957dcf269204462a27f7a2d', 'started_at': '2025-08-18T11:00:51.481624+00:00', 'ended_at': '2025-08-18T11:00:52.581610+00:00', 'span_data': {'type': 'response', 'response_id': 'resp-d4f5cbc0-91de-457e-9be3-627def83e53f'}, 'error': None}\n",
      "[Exporter] Export span: {'object': 'trace.span', 'id': 'span_0957dcf269204462a27f7a2d', 'trace_id': 'trace_04e821a8ac0744a083a1549fdd6c4465', 'parent_id': None, 'started_at': '2025-08-18T11:00:47.819925+00:00', 'ended_at': '2025-08-18T11:00:52.582087+00:00', 'span_data': {'type': 'agent', 'name': 'Doc QA', 'handoffs': [], 'tools': ['is_document_in_local_cache', 'convert_document_into_docling_document', 'export_docling_document_to_vector_db', 'search_documents'], 'output_type': 'str'}, 'error': None}\n",
      "[Exporter] Export trace_id=trace_8d9ec7fa83834e67a7a9c2a9dfbfc477, name=Agent workflow\n",
      "[Exporter] Export span: {'object': 'trace.span', 'id': 'span_bef67a9ee9e1435fa1ead60f', 'trace_id': 'trace_8d9ec7fa83834e67a7a9c2a9dfbfc477', 'parent_id': None, 'started_at': '2025-08-18T11:00:52.584203+00:00', 'ended_at': '2025-08-18T11:00:52.590072+00:00', 'span_data': {'type': 'mcp_tools', 'server': 'Docling MCP', 'result': ['is_document_in_local_cache', 'convert_document_into_docling_document', 'export_docling_document_to_vector_db', 'search_documents']}, 'error': None}\n",
      "[Exporter] Export span: {'object': 'trace.span', 'id': 'span_e1b42c22ff444548a27d546a', 'trace_id': 'trace_8d9ec7fa83834e67a7a9c2a9dfbfc477', 'parent_id': 'span_348fe1f320fc4919b1f12493', 'started_at': '2025-08-18T11:00:52.590319+00:00', 'ended_at': '2025-08-18T11:00:53.570078+00:00', 'span_data': {'type': 'response', 'response_id': 'resp-a311f00d-7cb6-4e17-a0fe-3e4148c623a8'}, 'error': None}\n",
      "[Exporter] Export span: {'object': 'trace.span', 'id': 'span_5579171ec0a946aa83e2ae1f', 'trace_id': 'trace_8d9ec7fa83834e67a7a9c2a9dfbfc477', 'parent_id': 'span_348fe1f320fc4919b1f12493', 'started_at': '2025-08-18T11:00:53.570201+00:00', 'ended_at': '2025-08-18T11:00:55.594609+00:00', 'span_data': {'type': 'function', 'name': 'search_documents', 'input': '{\"query\":\"How many pages were manually annotated in the dataset?\"}', 'output': '{\"type\":\"text\",\"text\":\"{\\\\n  \\\\\"answer\\\\\": \\\\\"91104 annotation instances were manually annotated across 7059 and 1591 pages respectively, totaling to 8259 pages with double annotations and 1631 pages with triple annotations. The exact count of unique pages cannot be determined from the given information.\\\\\"\\\\n}\",\"annotations\":null,\"meta\":null}', 'mcp_data': {'server': 'Docling MCP'}}, 'error': None}\n",
      "[Exporter] Export span: {'object': 'trace.span', 'id': 'span_88ecac551aa94d6d9b7a3b1a', 'trace_id': 'trace_8d9ec7fa83834e67a7a9c2a9dfbfc477', 'parent_id': 'span_348fe1f320fc4919b1f12493', 'started_at': '2025-08-18T11:00:55.594820+00:00', 'ended_at': '2025-08-18T11:00:55.600464+00:00', 'span_data': {'type': 'mcp_tools', 'server': 'Docling MCP', 'result': ['is_document_in_local_cache', 'convert_document_into_docling_document', 'export_docling_document_to_vector_db', 'search_documents']}, 'error': None}\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #008000; text-decoration-color: #008000\">Assistant&gt; The search query returned that in the dataset, a total of </span><span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">91104</span><span style=\"color: #008000; text-decoration-color: #008000\"> annotation instances were manually </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">annotated across </span><span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">7059</span><span style=\"color: #008000; text-decoration-color: #008000\"> and </span><span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">1591</span><span style=\"color: #008000; text-decoration-color: #008000\"> pages respectively. This means there are </span><span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">8259</span><span style=\"color: #008000; text-decoration-color: #008000\"> pages with double annotations and </span><span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">1631</span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">pages with triple annotations </span><span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">(</span><span style=\"color: #008000; text-decoration-color: #008000\">these numbers might have some issue to them possibly being truncated</span><span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">)</span><span style=\"color: #008000; text-decoration-color: #008000\">. </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">Unfortunately, it's not possible to determine exactly how many unique pages exist in the dataset based on the </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">available information.</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[32mAssistant> The search query returned that in the dataset, a total of \u001b[0m\u001b[1;32m91104\u001b[0m\u001b[32m annotation instances were manually \u001b[0m\n",
       "\u001b[32mannotated across \u001b[0m\u001b[1;32m7059\u001b[0m\u001b[32m and \u001b[0m\u001b[1;32m1591\u001b[0m\u001b[32m pages respectively. This means there are \u001b[0m\u001b[1;32m8259\u001b[0m\u001b[32m pages with double annotations and \u001b[0m\u001b[1;32m1631\u001b[0m\n",
       "\u001b[32mpages with triple annotations \u001b[0m\u001b[1;32m(\u001b[0m\u001b[32mthese numbers might have some issue to them possibly being truncated\u001b[0m\u001b[1;32m)\u001b[0m\u001b[32m. \u001b[0m\n",
       "\u001b[32mUnfortunately, it's not possible to determine exactly how many unique pages exist in the dataset based on the \u001b[0m\n",
       "\u001b[32mavailable information.\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Exporter] Export span: {'object': 'trace.span', 'id': 'span_6b46f0309e474aa2ad41861c', 'trace_id': 'trace_8d9ec7fa83834e67a7a9c2a9dfbfc477', 'parent_id': 'span_348fe1f320fc4919b1f12493', 'started_at': '2025-08-18T11:00:55.600716+00:00', 'ended_at': '2025-08-18T11:00:57.821127+00:00', 'span_data': {'type': 'response', 'response_id': 'resp-eabde418-a256-4d72-b066-a9134377a69f'}, 'error': None}\n",
      "[Exporter] Export span: {'object': 'trace.span', 'id': 'span_348fe1f320fc4919b1f12493', 'trace_id': 'trace_8d9ec7fa83834e67a7a9c2a9dfbfc477', 'parent_id': None, 'started_at': '2025-08-18T11:00:52.590122+00:00', 'ended_at': '2025-08-18T11:00:57.822614+00:00', 'span_data': {'type': 'agent', 'name': 'Doc QA', 'handoffs': [], 'tools': ['is_document_in_local_cache', 'convert_document_into_docling_document', 'export_docling_document_to_vector_db', 'search_documents'], 'output_type': 'str'}, 'error': None}\n"
     ]
    }
   ],
   "source": [
    "client = AsyncOpenAI(base_url=BASE_URL, api_key=API_KEY)\n",
    "\n",
    "# Configure the OpenAI provider that uses our AsyncOpenAI client for Llama Stack\n",
    "provider = OpenAIProvider(openai_client=client)\n",
    "\n",
    "# Tell OpenAI to dump traces to the console\n",
    "set_trace_processors([BatchTraceProcessor(exporter=ConsoleSpanExporter())])\n",
    "\n",
    "\n",
    "async def run_agent():\n",
    "    async with MCPServerStreamableHttp(\n",
    "        name=\"Docling MCP\",\n",
    "        params={\n",
    "            \"url\": \"http://localhost:8000/mcp\",\n",
    "            \"timeout\": 60.0,\n",
    "        },\n",
    "        client_session_timeout_seconds=60,\n",
    "    ) as server:\n",
    "        agent = Agent(\n",
    "            name=\"Doc QA\",\n",
    "            model=MODEL_ID,\n",
    "            instructions=\"You are a helpful assistant. Use the tools you have access to for providing relevant answers.\",\n",
    "            model_settings=ModelSettings(\n",
    "                temperature=1.0, top_p=0.9, tool_choice=\"required\"\n",
    "            ),\n",
    "            mcp_servers=[server],\n",
    "        )\n",
    "        session = SQLiteSession(str(uuid.uuid4()))\n",
    "        print(f\"Created session_id={session.session_id} for Agent({agent.name})\")\n",
    "\n",
    "        user_prompts = [\n",
    "            \"Convert the PDF document on https://arxiv.org/pdf/2206.01062 to DoclingDocument.\",\n",
    "            \"Export the document to the vectordb.\",\n",
    "            \"Search in the document: How many pages were manually annotated in the dataset?\",\n",
    "        ]\n",
    "        for prompt in user_prompts:\n",
    "            console.print(f\"[cyan]User> {prompt}[/cyan]\")\n",
    "            response = await Runner.run(\n",
    "                agent,\n",
    "                prompt,\n",
    "                session=session,\n",
    "                run_config=RunConfig(model_provider=provider),\n",
    "            )\n",
    "            console.print(f\"[green]Assistant> {response.final_output}[/green]\")\n",
    "\n",
    "\n",
    "await run_agent()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama-index",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
